from ultralytics import YOLO
# 加载模型
# model = YOLO("runs/detect/train21/weights/best_modified.pt")  # 推荐加载预训练权重
model = YOLO("weights/yolov8m.pt")  # 推荐加载预训练权重
if __name__ == '__main__':
    # 训练配置
    results = model.train(
        data="ultralytics/cfg/datasets/helmet3.yaml",
        epochs=500,
        lr0=0.01,  # 减小学习率防止破坏原有权重
        # resume=True,  # 关键参数
        batch=16,
        imgsz=640,
        amp=True,
        device=0,
        project="runs/detect",
        workers=0,  # 关键！Windows下禁用多进程数据加载
        rect=True,  # 关键！启用矩形训练（保持原比例）
        augment=True,
        hsv_h=0.015,  # 增强色相扰动
        copy_paste=0.4,  # 小目标复制增强
        erasing=0.7,  # 随机擦除增强
        degrees=10.0  # 旋转增强提升至10度
    )

# yolo task=detect mode=train model=weights/yolov8n.pt data=ultralytics/cfg/datasets/CCPD-test.yaml batch=16 epochs=50 imgsz=416 workers=1 device=0 project=runs/detect